package com.nbsaas.boot.lock.impl;

import com.nbsaas.boot.lock.api.LockService;

import java.util.concurrent.*;
import java.util.concurrent.locks.*;

public class StampedLockService implements LockService {

    private final ConcurrentMap<String, StampedLock> lockMap = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, Long> lockStamps = new ConcurrentHashMap<>(); // 保存锁的stamp

    @Override
    public boolean lock(String key) {
        StampedLock lock = lockMap.computeIfAbsent(key, k -> new StampedLock());
        long stamp = lock.tryWriteLock(); // 非阻塞获取写锁
        if (stamp != 0) {
            lockStamps.put(key, stamp); // 保存获取的锁的stamp
        }
        return stamp != 0;
    }

    @Override
    public boolean lock(String key, long timeout, TimeUnit unit) {
        StampedLock stampedLock = lockMap.computeIfAbsent(key, k -> new StampedLock());
        long timeoutMillis = unit.toMillis(timeout);
        long startTime = System.currentTimeMillis();
        
        while (System.currentTimeMillis() - startTime < timeoutMillis) {
            long stamp = stampedLock.writeLock();
            if (stamp != 0) {
                lockStamps.put(key, stamp); // 保存获取的锁的stamp
                return true;
            }
        }
        return false;  // 超时未获取到锁
    }

    @Override
    public boolean unlock(String key) {
        StampedLock stampedLock = lockMap.get(key);
        if (stampedLock != null) {
            Long stamp = lockStamps.get(key);
            if (stamp != null) {
                stampedLock.unlockWrite(stamp);  // 使用正确的stamp释放锁
                lockStamps.remove(key);  // 清除该key对应的stamp
                return true;
            }
        }
        return false;
    }

    @Override
    public boolean isLocked(String key) {
        StampedLock stampedLock = lockMap.get(key);
        return stampedLock != null && stampedLock.isWriteLocked();
    }
}
